# -*- coding: utf-8 -*-
"""
Created on Thu Jul 30 18:08:24 2020

@author: 1
"""

import torch

def get_network(args, device):
    """ return specific network according to the args.net
    """
    net = None
    netname = ''

    if args.net == 'resnet50':
        from model.resnet50 import resnet50  # 引入模型架构的代码  model/xxx.py
        net = resnet50()
        netname = 'ResNet50'
    
    elif args.net == 'resnet18':
        netname = 'ResNet18'
        pass

    elif args.net == 'resnet101':
        netname = 'ResNet101'
        pass

    net.to(device)

    return net , netname + device.split(':')[0]


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('-gpu', type=bool, default=True, help='use gpu or not')
    parser.add_argument('-gpunum', type=int, default=0, help='use the kth gpu')
    parser.add_argument('-net', type=str, default='resnet50', help='network type')
    args = parser.parse_args()

    net, netname = get_network(args)
    print(netname)
    # print(net)
